package cn.baddonkey.photography.agent.config;

import cn.baddonkey.photography.agent.graph.NodeStatus;
import cn.baddonkey.photography.agent.node.ImageHandlerNode;
import cn.baddonkey.photography.agent.node.ImageIntroNode;
import cn.baddonkey.photography.agent.node.ImageJudgeNode;
import com.alibaba.cloud.ai.graph.GraphRepresentation;
import com.alibaba.cloud.ai.graph.KeyStrategy;
import com.alibaba.cloud.ai.graph.KeyStrategyFactory;
import com.alibaba.cloud.ai.graph.OverAllState;
import com.alibaba.cloud.ai.graph.StateGraph;
import com.alibaba.cloud.ai.graph.action.AsyncEdgeAction;
import com.alibaba.cloud.ai.graph.action.NodeAction;
import com.alibaba.cloud.ai.graph.exception.GraphStateException;
import com.alibaba.cloud.ai.graph.state.strategy.ReplaceStrategy;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;

import java.util.HashMap;
import java.util.Map;

import static com.alibaba.cloud.ai.graph.action.AsyncNodeAction.node_async;

/**
 * @author yingzi
 * @since 2025/8/26
 */
@Configuration
public class GraphConfiguration {

    private static final Logger logger = LoggerFactory.getLogger(GraphConfiguration.class);

    @Bean
    public StateGraph parallelStreamGraph() throws GraphStateException {
        KeyStrategyFactory keyStrategyFactory = () -> {
            HashMap<String, KeyStrategy> keyStrategyHashMap = new HashMap<>();

            // 用户输入
            keyStrategyHashMap.put("imageUrl", new ReplaceStrategy());
            keyStrategyHashMap.put("ImageHandlerNodecontentText", new ReplaceStrategy());
            keyStrategyHashMap.put("ImageHandlerNodecontent", new ReplaceStrategy());

            keyStrategyHashMap.put("expandercontent", new ReplaceStrategy());
            keyStrategyHashMap.put("expandercontentText", new ReplaceStrategy());
            keyStrategyHashMap.put("threadId", new ReplaceStrategy());

            keyStrategyHashMap.put("ImageJudgeNode", new ReplaceStrategy());

            return keyStrategyHashMap;
        };

        Map<String, NodeStatus> node2Status = new HashMap<>();

        StateGraph stateGraph = new StateGraph(keyStrategyFactory)
                .addNode("ImageIntroNode", node_async(new ImageIntroNode(node2Status)))
                .addNode("ImageHandlerNode", node_async(new ImageHandlerNode(node2Status)))
                .addNode("ImageJudgeNode", node_async(new ImageJudgeNode(node2Status)))

                .addEdge(StateGraph.START, "ImageIntroNode")
                .addEdge("ImageIntroNode", "ImageJudgeNode")
                .addConditionalEdges("ImageJudgeNode", AsyncEdgeAction.edge_async(state -> {
                    Boolean value = state.value("ImageJudgeNode", false);
                    return value ? "ImageHandlerNode" : StateGraph.END;
                }), Map.of("ImageHandlerNode", "ImageHandlerNode", StateGraph.END, StateGraph.END))
                .addEdge("ImageHandlerNode", StateGraph.END);

        // 添加 PlantUML 打印
        GraphRepresentation representation = stateGraph.getGraph(GraphRepresentation.Type.PLANTUML,
                "expander flow");
        logger.info("\n=== expander UML Flow ===");
        logger.info(representation.content());
        logger.info("==================================\n");

        return stateGraph;
    }

}